# 路径token化

import os
import torch
from src.configs import base_config
from src.data_preprocessing.svg_constants import SVG_COMMANDS

class PathTokenizer:
    def __init__(self):
        # 使用svg_constants.py中的标准命令定义
        self.vocab = {
            cmd: idx for idx, cmd in enumerate(SVG_COMMANDS)
        }
        # 添加坐标token
        self.vocab.update({
            f'coord_{i}': len(SVG_COMMANDS) + i 
            for i in range(96)
        })
    
    def tokenize(self, svg_content):
        """增强版SVG路径token化方法"""
        # 提取纯路径数据
        if '<path' in svg_content:
            svg_content = svg_content.split('d="')[-1].split('"')[0]
        
        tokens = []
        for cmd in svg_content.split():
            if cmd in self.vocab:  # 检查是否是已知命令
                tokens.append(self.vocab[cmd])
            else:
                try:
                    # 尝试解析为坐标值
                    coord = int(float(cmd) * 10) % 96
                    tokens.append(self.vocab[f'coord_{coord}'])
                except ValueError:
                    continue  # 跳过无法解析的字符串
        return torch.tensor(tokens, dtype=torch.long)

def main():
    import argparse
    parser = argparse.ArgumentParser(description='SVG路径token化工具')
    parser.add_argument("--svg_root", default="/Users/weiyin/Workspaces/TypusMotrix/data/processed/svg_paths")
    parser.add_argument("--output_root", default="/Users/weiyin/Workspaces/TypusMotrix/data/processed/tokens")
    args = parser.parse_args()

    tokenizer = PathTokenizer()
    
    # 处理所有字体子目录
    font_dirs = [d for d in os.listdir(args.svg_root) 
                if os.path.isdir(os.path.join(args.svg_root, d))]
    
    for font_dir in font_dirs:
        font_path = os.path.join(args.svg_root, font_dir)
        output_dir = os.path.join(args.output_root, font_dir)
        os.makedirs(output_dir, exist_ok=True)

        print(f"正在处理字体目录: {font_dir}")
        for svg_file in os.listdir(font_path):
            if svg_file.endswith('.svg'):
                svg_path = os.path.join(font_path, svg_file)
                with open(svg_path) as f:
                    tokens = tokenizer.tokenize(f.read())
                
                torch.save(
                    tokens,
                    os.path.join(output_dir, f"{os.path.splitext(svg_file)[0]}.pt")
                )

if __name__ == "__main__":
    main()